import pyqtgraph as pg
import numpy as np
from pyqtgraph.Qt import QtCore
import pyaudio

###################################
# 音频文件的PSD功率谱密度分析-实时模式
# 从声卡读取流式音频数据，生成PSD
###################################
#            参 数 
# 采样率：48KHz
# 时间窗口大小：1s (48000 points)
# step: 0.1s (4800 points)# #
###################################


CHUNK = 4800                 # samples per frame
FORMAT = pyaudio.paInt16     # audio format (bytes per sample?)
CHANNELS = 1                 # single channel for microphone
RATE = 48000                 # samples per second
pya = pyaudio.PyAudio()
count = 0
buffer = np.zeros(48000, dtype=np.int16)

frequency_bin_min = 9900
frequency_bin_max = 10100

app = pg.QtGui.QApplication([])
win = pg.GraphicsWindow(title="从麦克风采集声音并进行快速傅里叶变换")
win.resize(1000,500)  #设置窗口大小
p = win.addPlot()
data = np.random.random(size=50)
curve = p.plot(data)  #在坐标p中绘图并返回图形对象    
p.setYRange(0,15,padding = 0)


stream = pya.open(
            format=FORMAT,
            channels=CHANNELS,
            rate=RATE,
            input=True,
            output=True,
            frames_per_buffer=CHUNK
        )

def update():
    global data, curve, buffer, count
    s_data = stream.read(CHUNK, exception_on_overflow=False)
    wave_data = np.frombuffer(s_data, dtype=np.short)
    if count<10:
        buffer[(count)*CHUNK:(count+1)*CHUNK] = wave_data
    else:
        buffer[0:43200] = buffer[4800:48000]
        buffer[43200:48000] = wave_data[0:4800]

    _r = np.log10(np.abs(np.fft.rfft(buffer))**2/len(buffer))
    output = _r[frequency_bin_min:frequency_bin_max]
    curve.setData(output)

    count += 1   

timer = QtCore.QTimer()  #创建一个定时器
timer.timeout.connect(update)  #设置定时器执行函数
timer.start(1)  #启动
#参数：间隔毫秒数

app.exec_()